#include "crfTrainImp0.h"
#include "lbfgs.h"
//////////////////////////////////////////////////////////////////////////

crfTrainImp0::crfTrainImp0()
{
	_feature_num_all=0;
	_depth=5;
	_prior=0;
	_sen_num = 0;
	_sigma=10;
	_temp_space.set_size(80*1024);
	_eta = 0.005;

}

void crfTrainImp0::Train(char *fea_file,char* lambda_file,vector<string> trainfiles,int save_step,int save_startid,bool bcheck)
{
	char lambda_save_path[512];
	strncpy(lambda_save_path,lambda_file,512);
	char* tfvvv=strrchr(lambda_save_path,'\\');
	if(tfvvv)
		*tfvvv=0;
	else if (tfvvv=strrchr(lambda_save_path,'/'))
		*tfvvv=0;
	else
		strncpy(lambda_save_path,"./lambda",512);
	//////////////////////////////////////////////////////////////////////////
	TempleteFea* _temfea = new TempleteFea(1,10);
	_temfea->load_model(fea_file,true);
	_feature_num_all= _temfea->_feature_num;
	_lambda = new double[_feature_num_all];
	_gradient = new double[_feature_num_all];
	memset(_lambda,0,_feature_num_all*sizeof(double));
	memset(_gradient,0,_feature_num_all*sizeof(double));
	ReadBinDouble(lambda_file,_feature_num_all,_lambda);
	//----------------------------------------------------------------
	_foMst._tag_size = _temfea->_tag2idx.size()+1;
	_foMst._lambda = _lambda;
	_foMst._gradient = _gradient;
	int begin_id = _temfea->get_label_index((string)"0o");
	assert_loc(begin_id!=-1);
	_foMst._tag_sentence_begin = begin_id+1;
	int _lambda_size = _feature_num_all;
	int _work_bytes=_feature_num_all*2*sizeof(double);
	pool_work_space.alloc(_work_bytes);
	cout<<"allocating "<<_work_bytes<<" bytes"<<endl;
	char* _work_space=(char*)pool_work_space;
	cout<<_work_bytes<<" bytes allocated"<<endl;
	//----------------------------------------------------------------
	double *w0,*w1;
	if(_prior==0){
		w0=NULL;
		w1=NULL;
	}else{
		w0=(double *)_work_space;
		w1=((double *)_work_space)+_feature_num_all;
	}

	LBFGS l;
	bool init_success=l.init(_feature_num_all,_depth,_prior);
	if(!init_success){
		if(_prior==0)
			printf("can not allocate enough memory, use option -m 1 to reduce memory requirement\n");
		else
			printf("can not allocate enough memory, use option -a 1 or -a 2 to train\n");
		exit(1);
	}
	//////////////////////////////////////////////////////////////////////////
	double old_fvalue=1e+37;
	char disp_buff[1000];
	int converge=0;

	vector<sentence> sen_poolck;
	memPool<char> vec_pools;
	char path[512];
	int start_train_id = bcheck? 1:0;
	vector<sentence> sen_pooltrain;
	memPool<char> vec_poolstrain;
	if(bcheck)
		read_senbin((char*)trainfiles[0].c_str(),sen_pooltrain,vec_poolstrain);

	for (int epc = start_train_id;epc<10000;epc++)
	{
		memset(_gradient,0,sizeof(double)*_lambda_size);
		double fvalue =0;

		printf("\n");
		int error_num = 0;
		int total_num = 0;
		for (int ibatch=0;ibatch<trainfiles.size();ibatch++)
		{
			printf("@begin handle %d\n",ibatch);
			read_senbin((char*)trainfiles[ibatch].c_str(),sen_poolck,vec_pools);
			//-------
			for (unsigned int isen=0;isen<sen_poolck.size();isen++) {
				fvalue += _foMst.process_sen(sen_poolck[isen]);
				cout<<"done["<<isen<<"]job|len="<<sen_poolck[isen].l<<"["<<fvalue<<"]...";
			}
			//-------
			printf("\n@end handle %d>>>\n",ibatch);
			//-------
			sen_poolck.clear();
			//-------
		}

		if(bcheck) {
			printf("\n@begin check\n");
			for (unsigned int isen=0;isen<sen_pooltrain.size();isen++) {
				vector<vector<int>> labels;
				vector< vector<double> >predict_prob;
				_foMst.process_decode(sen_pooltrain[isen],labels,predict_prob);
				for(int i=0;i<labels.size();i++)
				{
					if(labels[i][0]!=sen_pooltrain[isen].labels[i])
						error_num++;
					total_num++;
				}
			}
			printf("train error ratio is %f\n",error_num/(float)total_num);
		}

		for(int j=0;j<_lambda_size;j++){
			fvalue+=_lambda[j] * _lambda[j] / (2*_sigma);
			_gradient[j]+=_lambda[j]/_sigma;
		}

		double diff= fabs((old_fvalue-fvalue)/old_fvalue) ;
		int selected=0;
		for(int j=0;j<_lambda_size;j++)
			if(_lambda[j]!=0)
				selected++;
		sprintf(disp_buff,"\niter: %d act: %d diff: %lf obj: %lf\n", epc+save_startid, selected, diff, fvalue);
		cout<<disp_buff<<endl;
		old_fvalue=fvalue;
		if(diff<_eta)
			converge++;
		else
			converge=0;
		if(converge==3) {//success
			sprintf(path,"%s/lambdaOrg0_ok",lambda_save_path);
			SaveTxtDouble(path,_feature_num_all,_lambda);
			sprintf(path,"%s/lambdaOrg0_okbin",lambda_save_path);
			SaveBinDouble(path,_feature_num_all,_lambda);
			break;
		}

		int ret=l.optimize(_lambda,&fvalue,_gradient,0,w0,w1);
		if(ret<0){
			cout<<"lbfgs error"<<endl;
			break;
		}else if(ret==0){
			break;//success
		}

		if( !(epc % save_step) ) {
			sprintf(path,"%s/lambdaOrg0[%d]",lambda_save_path,epc+save_startid);
			SaveTxtDouble(path,_feature_num_all,_lambda);
			sprintf(path,"%s/lambdaOrg0[%d]_bin",lambda_save_path,epc+save_startid);
			SaveBinDouble(path,_feature_num_all,_lambda);
		}

		if( !(epc % save_step) ) {
			sprintf(path,"%s/lambdaOrg0[%d]_gradient",lambda_save_path,epc+save_startid);
			FILE *fp = fopen(path, "w");
			if (!fp) return ;
			for (int i=0;i<_feature_num_all;i++)
				fprintf(fp, "%lf\n",_gradient[i]);
			fclose(fp);

		}

		cout<<endl<<"epoch="<<epc+save_startid<<"!"<<endl;
	}
	cout<<"***finish I-O job!"<<endl;
	return ;
}
